import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, models, transforms
from torch.utils.data import DataLoader
from tqdm import tqdm
import time
import os
import datetime

def main():
    # -----------------------------
    # 1️⃣ 环境与设备检测
    # -----------------------------
    device = (
        torch.device("mps") if torch.backends.mps.is_available()
        else torch.device("cuda") if torch.cuda.is_available()
        else torch.device("cpu")
    )
    print(f"📌 当前设备: {device}")
    print(f"🔍 PyTorch 版本: {torch.__version__}")
    print(f"🔋 是否支持 MPS: {torch.backends.mps.is_available()}")
    print(f"🔋 是否支持 CUDA: {torch.cuda.is_available()}")

    # -----------------------------
    # 2️⃣ 数据预处理与加载
    # -----------------------------
    transform = transforms.Compose([
        transforms.Resize((128, 128)),   # 调整图片尺寸为 128x128
        transforms.ToTensor(),           # 转为张量
    ])

    train_dataset = datasets.ImageFolder("data/processed/train", transform=transform)
    val_dataset = datasets.ImageFolder("data/processed/val", transform=transform)

    batch_size = 32
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2)
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=2)

    print(f"📦 数据加载完成 -> 训练样本: {len(train_dataset)}, 验证样本: {len(val_dataset)}")

    # -----------------------------
    # 3️⃣ 模型构建
    # -----------------------------
    model = models.resnet18(weights="IMAGENET1K_V1")
    model.fc = nn.Linear(model.fc.in_features, 2)  # 猫/狗二分类
    model = model.to(device)

    # -----------------------------
    # 4️⃣ 损失函数与优化器
    # -----------------------------
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5)  # 每 10 轮减半学习率

    # -----------------------------
    # 5️⃣ 训练配置
    # -----------------------------
    num_epochs = 30  # 🔁 合理的高精度训练批次
    train_losses = []
    val_accuracies = []
    start_time = time.time()

    # -----------------------------
    # 6️⃣ 训练与验证
    # -----------------------------
    try:
        for epoch in range(num_epochs):
            print(f"\n🚀 [Epoch {epoch+1}/{num_epochs}] 开始训练...")
            model.train()
            running_loss = 0.0

            # tqdm 进度条
            progress_bar = tqdm(train_loader, desc=f"Epoch {epoch+1}", ncols=100)
            for images, labels in progress_bar:
                images, labels = images.to(device), labels.to(device)

                optimizer.zero_grad()
                outputs = model(images)
                loss = criterion(outputs, labels)
                loss.backward()
                optimizer.step()

                running_loss += loss.item()
                progress_bar.set_postfix({"loss": f"{loss.item():.4f}"})

            avg_loss = running_loss / len(train_loader)
            train_losses.append(avg_loss)
            print(f"📉 平均训练损失: {avg_loss:.4f}")

            # 验证阶段
            model.eval()
            correct, total = 0, 0
            with torch.no_grad():
                for images, labels in val_loader:
                    images, labels = images.to(device), labels.to(device)
                    outputs = model(images)
                    _, preds = torch.max(outputs, 1)
                    total += labels.size(0)
                    correct += (preds == labels).sum().item()

            acc = 100 * correct / total
            val_accuracies.append(acc)
            print(f"🎯 验证准确率: {acc:.2f}%")

            scheduler.step()
            print(f"🧠 当前学习率: {optimizer.param_groups[0]['lr']:.6f}")

    except KeyboardInterrupt:
        print("\n⚠️ 手动中断训练，正在保存模型...")

    finally:
        # -----------------------------
        # 7️⃣ 保存模型
        # -----------------------------
        output_path = "cat_dog_model_mps.pth"
        torch.save(model.state_dict(), output_path)
        total_time = time.time() - start_time
        print("\n✅ 训练完成")
        print(f"📁 模型已保存为: {output_path}")
        print(f"⏱️ 总耗时: {total_time/60:.2f} 分钟")

if __name__ == "__main__":
    main()